import llparse
import arbb

import sys

def isFormalVariable(fun, name):

    for (n, arg) in enumerate(fun.args):
        (ty, argname) = arg
        if name == argname:
            return (True, n)

    return (False, -1)

def escapeRegName(s):

    ss = '_' + s[1:]

    return ss.replace('.', '_')

def convertType(ty):

    tymap = {'float' : 'f32'}
    
    return tymap[ty]

def convertVar(s):

    ss = 'var_' + s[1:]

    return ss.replace('.', '_')

def emitConstInst(ctx, val):

    i = arbb.Const(ctx, val)

    return i

def convertRetOp(ctx, fun, inst):

    ii = []

    if inst.dst == None: # return void
        return ii

    print "---------- ret"
    dst = "var_ret"
    src = convertVar(inst.dst)
    isout = 1
    ii.append(arbb.GetParam(fun, dst, isout, 0))

    ii.append(arbb.AssignOp(fun, dst, src))

    return ii
    

def convertLoadOp(ctx, fun, inst):

    ii = []

    src = convertVar(inst.srcs[0])
    ty  = inst.srcs[1]
    dst = convertVar(inst.dst)
    name = escapeRegName(inst.dst)

    ii.append(arbb.LocalVariable(fun, arbb.ScalarTy(ctx, convertType(ty)), name))

    ii.append(arbb.AssignOp(fun, dst, src))

    return ii

def convertStoreOp(ctx, fun, fn, inst):

    ii = []

    src = convertVar(inst.srcs[0])
    dst = convertVar(inst.dst)

    (ret, n) = isFormalVariable(fn, inst.srcs[0])
    if ret == True:
        isout = 0
        ii.append(arbb.GetParam(fun, src, isout, n))

    ii.append(arbb.AssignOp(fun, dst, src))

    return ii

def convertBinOp(ctx, fun, inst):

    ii = []

    opmap = {'fadd' : 'add',
             'fsub' : 'sub',
             'fmul' : 'mul',
             'fdiv' : 'div'}
    
    arbbop = opmap[inst.op]
    
    if type(inst.srcs[0]) == str:
        src0 = convertVar(inst.srcs[0])
    else:   # fp const val
        i = emitConstInst(ctx, inst.srcs[0])
        src0 = i.name
        ii.append(i)
        
    if type(inst.srcs[1]) == str:
        src1 = convertVar(inst.srcs[1])
    else:   # fp const val
        j = emitConstInst(ctx, inst.srcs[1])
        src1 = j.name
        ii.append(j)

    dst = convertVar(inst.dst)
    name = escapeRegName(inst.dst)
    ty = inst.srcs[2]

    # Assume f32
    ii.append(arbb.LocalVariable(fun, arbb.ScalarTy(ctx, convertType(ty)), name))
    ii.append(arbb.BinOp(fun, arbbop, dst, src0, src1))

    return ii
    
def convertInst(ctx, fun, fn, inst):

    ii = []
    binops = ['fadd', 'fsub', 'fmul', 'fdiv']

    if inst.op == 'alloca':
        ty = inst.srcs[0]
        print inst.dst
        name = escapeRegName(inst.dst)
        i = arbb.LocalVariable(fun, arbb.ScalarTy(ctx, convertType(ty)), name)
        ii.append(i)
    elif inst.op == 'load':
        ii += convertLoadOp(ctx, fun, inst)
    elif inst.op == 'store':
        ii += convertStoreOp(ctx, fun, fn, inst)
    elif inst.op == 'ret':
        ii += convertRetOp(ctx, fun, inst)
    elif inst.op in binops:
        ii += convertBinOp(ctx, fun, inst)

    return ii

def convertFunction(ctx, fn):

    print fn.name

    outTys = []
    inTys  = []

    outTys.append(arbb.ScalarTy(ctx, convertType(fn.resTy)))

    for arg in fn.args:
        (ty, name) = arg
        inTys.append(arbb.ScalarTy(ctx, convertType(ty)))
        
    fun = arbb.Function(ctx, fn.name, outTys, inTys)

    for inst in fn.insts:
        ii = convertInst(ctx, fun, fn, inst)
        if ii is not None :
            fun.insts += ii

    print fun.emit()
    return fun

def main():

    if len(sys.argv) < 2:
        print "Usage: conv input.ll"
        sys.exit(1)

    fns = llparse.parseLLVMIR(sys.argv[1])

    ctx = arbb.Context()

    print "=====------===="
    arbbfns = []
    for fn in fns:
        arbbfn = convertFunction(ctx, fn)
        if arbbfn != None:
            arbbfns.append(arbbfn)
    print "=====------===="
    
    arbb.emitCode(ctx, arbbfns)

if __name__ == '__main__':
    main()
